/*
 * 虚拟内存管理实现
 * 实现Sv39页表管理系统
 * 参考xv6的实现，适配riscv-os项目
 */

#include "vm.h"
#include "kalloc.h"
#include "../fs/types.h"
#include "riscv.h"
#include "memlayout.h"
#include "../fs/printf.h"

// 内核页表
pagetable_t kernel_pagetable;

// 内存设置和复制
static void memset(void *dst, int c, uint64 n)
{
    char *cdst = (char*)dst;
    for (uint64 i = 0; i < n; i++) {
        cdst[i] = c;
    }
}

/*
 * 创建一个新的页表
 * 返回: 新页表的指针，失败时返回0
 */
pagetable_t create_pagetable(void)
{
    pagetable_t pagetable;
    
    pagetable = (pagetable_t)alloc_page();
    if (pagetable == 0) {
        return 0;
    }
    
    // 清零页表内容
    memset(pagetable, 0, PGSIZE);
    return pagetable;
}

/*
 * 页表遍历函数
 * 返回虚拟地址va对应的页表项指针
 * pagetable: 页表根指针
 * va: 虚拟地址
 * alloc: 是否在需要时分配中间页表
 */
pte_t* walk(pagetable_t pagetable, uint64 va, int alloc)
{
    // 检查虚拟地址是否超出范围
    if (va >= MAXVA) {
        printf("walk: virtual address too large: 0x%lx\n", va);
        return 0;
    }
    
    // 从第2级开始向下遍历（Sv39有3级页表：2, 1, 0）
    for (int level = 2; level > 0; level--) {
        pte_t *pte = &pagetable[PX(level, va)];
        
        if (*pte & PTE_V) {
            // 页表项有效，获取下一级页表
            pagetable = (pagetable_t)PTE2PA(*pte);
        } else {
            // 页表项无效
            if (!alloc) {
                return 0;  // 不分配新页表，返回失败
            }
            
            // 分配新的页表页
            pagetable = (pagetable_t)alloc_page();
            if (pagetable == 0) {
                return 0;  // 内存不足
            }
            
            // 清零新页表
            memset(pagetable, 0, PGSIZE);
            
            // 设置页表项指向新页表
            *pte = PA2PTE(pagetable) | PTE_V;
        }
    }
    
    // 返回第0级页表项的指针
    return &pagetable[PX(0, va)];
}

/*
 * 虚拟地址转物理地址
 * pagetable: 页表
 * va: 虚拟地址
 * 返回: 物理地址，失败时返回0
 */
uint64 walkaddr(pagetable_t pagetable, uint64 va)
{
    pte_t *pte;
    uint64 pa;
    
    if (va >= MAXVA) {
        return 0;
    }
    
    pte = walk(pagetable, va, 0);
    if (pte == 0) {
        return 0;
    }
    
    if ((*pte & PTE_V) == 0) {
        return 0;
    }
    
    pa = PTE2PA(*pte);
    return pa;
}

/*
 * 建立虚拟地址到物理地址的映射
 * pagetable: 页表
 * va: 虚拟地址起始
 * size: 映射大小
 * pa: 物理地址起始
 * perm: 权限位
 * 返回: 成功返回0，失败返回-1
 */
int map_pages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
{
    uint64 a, last;
    pte_t *pte;
    
    // 检查地址对齐
    if ((va % PGSIZE) != 0) {
        printf("map_pages: va not aligned: 0x%lx\n", va);
        return -1;
    }
    
    if ((size % PGSIZE) != 0) {
        printf("map_pages: size not aligned: %ld\n", size);
        return -1;
    }
    
    if (size == 0) {
        printf("map_pages: size is zero\n");
        return -1;
    }
    
    a = va;
    last = va + size - PGSIZE;
    
    for (;;) {
        pte = walk(pagetable, a, 1);
        if (pte == 0) {
            printf("map_pages: walk failed for va 0x%lx\n", a);
            return -1;
        }
        
        if (*pte & PTE_V) {
            printf("map_pages: remap attempt at va 0x%lx\n", a);
            return -1;
        }
        
        *pte = PA2PTE(pa) | perm | PTE_V;
        
        if (a == last) {
            break;
        }
        
        a += PGSIZE;
        pa += PGSIZE;
    }
    
    return 0;
}

/*
 * 在内核页表中建立映射
 * kpgtbl: 内核页表
 * va: 虚拟地址
 * pa: 物理地址
 * sz: 大小
 * perm: 权限
 */
void kvmmap(pagetable_t kpgtbl, uint64 va, uint64 pa, uint64 sz, int perm)
{
    if (map_pages(kpgtbl, va, sz, pa, perm) != 0) {
        printf("kvmmap failed: va=0x%lx, pa=0x%lx, sz=%ld\n", va, pa, sz);
    }
}

/*
 * 创建内核页表
 * 返回: 内核页表指针
 */
static pagetable_t kvmmake(void)
{
    pagetable_t kpgtbl;
    
    kpgtbl = create_pagetable();
    if (kpgtbl == 0) {
        printf("kvmmake: failed to allocate kernel page table\n");
        return 0;
    }
    
    // 映射UART寄存器
    kvmmap(kpgtbl, UART0, UART0, PGSIZE, PTE_R | PTE_W);
    
    // 映射VirtIO设备
    kvmmap(kpgtbl, VIRTIO0, VIRTIO0, PGSIZE, PTE_R | PTE_W);
    
    // 映射PLIC（平台级中断控制器）
    kvmmap(kpgtbl, PLIC, PLIC, 0x4000000, PTE_R | PTE_W);
    
    // 映射内核代码段（只读+可执行）
    extern char etext[];  // 内核代码结束地址
    kvmmap(kpgtbl, KERNBASE, KERNBASE, (uint64)etext - KERNBASE, PTE_R | PTE_X);
    
    // 映射内核数据段和物理内存（读写）
    kvmmap(kpgtbl, (uint64)etext, (uint64)etext, PHYSTOP - (uint64)etext, PTE_R | PTE_W);
    
    return kpgtbl;
}

/*
 * 初始化内核虚拟内存
 */
void kvminit(void)
{
    printf("Initializing kernel virtual memory...\n");
    
    kernel_pagetable = kvmmake();
    if (kernel_pagetable == 0) {
        printf("kvminit: failed to create kernel page table\n");
        return;
    }
    
    printf("Kernel page table created at 0x%lx\n", (uint64)kernel_pagetable);
}

/*
 * 激活内核页表（在每个hart上调用）
 */
void kvminithart(void)
{
    // 设置SATP寄存器启用Sv39分页
    w_satp(MAKE_SATP(kernel_pagetable));
    
    // 刷新TLB
    sfence_vma();
    
    printf("Kernel page table activated on hart\n");
}

/*
 * 创建用户页表
 */
pagetable_t uvmcreate(void)
{
    pagetable_t pagetable;
    
    pagetable = create_pagetable();
    if (pagetable == 0) {
        return 0;
    }
    
    return pagetable;
}

/*
 * 递归释放页表
 */
static void freewalk(pagetable_t pagetable, int level)
{
    // 遍历页表的所有条目
    for (int i = 0; i < 512; i++) {
        pte_t pte = pagetable[i];
        if (pte & PTE_V) {
            if (level > 0) {
                // 中间级页表，递归释放
                uint64 child = PTE2PA(pte);
                freewalk((pagetable_t)child, level - 1);
                pagetable[i] = 0;
            }
        }
    }
    free_page((void*)pagetable);
}

/*
 * 销毁页表
 */
void destroy_pagetable(pagetable_t pt)
{
    freewalk(pt, 2);
}

/*
 * 打印页表内容（调试用）
 */
void dump_pagetable(pagetable_t pt, int level)
{
    for (int i = 0; i < 512; i++) {
        pte_t pte = pt[i];
        if (pte & PTE_V) {
            for (int j = 0; j < level; j++) {
                printf("  ");
            }
            printf("Level %d[%d]: PTE=0x%lx PA=0x%lx", level, i, pte, PTE2PA(pte));
            if (pte & PTE_R) printf(" R");
            if (pte & PTE_W) printf(" W");
            if (pte & PTE_X) printf(" X");
            if (pte & PTE_U) printf(" U");
            printf("\n");
            
            if (level > 0) {
                dump_pagetable((pagetable_t)PTE2PA(pte), level - 1);
            }
        }
    }
}